{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "TUQqwe_N6J-w"
      },
      "source": [
        "Copyright 2018 The Dopamine Authors.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\"); you may not use this file except in compliance with the License. You may obtain a copy of the License at\n",
        "\n",
        "https://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "y08iF10S6N-6"
      },
      "source": [
        "# Dopamine Agent visualizer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "2cn16cOg6VLZ"
      },
      "source": [
        "This colaboratory demonstrates how to use the agent visualizer functionality in Dopamine. It uses a pre-trained Rainbow agent on SpaceInvaders\n",
        "and generates the video over 1000 steps of agent play.\n",
        "\n",
        "Note that it will save all the files to a temp directory in your runtime.\n",
        "\n",
        "Then run all the cells in order."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "0uZ5ifEc6DKa"
      },
      "outputs": [],
      "source": [
        "# @title Install necessary packages.\n",
        "!pip install -U dopamine-rl"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "Q7BnRy2w6xEK"
      },
      "outputs": [],
      "source": [
        "# @title Download an example checkpoint (Rainbow on SpaceInvaders)\n",
        "!gsutil -q -m cp -R gs://download-dopamine-rl/colab/samples/rainbow/SpaceInvaders_v4/checkpoints/tf_ckpt-199.data-00000-of-00001 /tmp\n",
        "!gsutil -q -m cp -R gs://download-dopamine-rl/colab/samples/rainbow/SpaceInvaders_v4/checkpoints/tf_ckpt-199.index /tmp\n",
        "!gsutil -q -m cp -R gs://download-dopamine-rl/colab/samples/rainbow/SpaceInvaders_v4/checkpoints/tf_ckpt-199.meta /tmp"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "both",
        "colab": {},
        "colab_type": "code",
        "id": "kpomDnVn9Ig1"
      },
      "outputs": [],
      "source": [
        "# @title Generate the video\n",
        "from dopamine.utils import example_viz_lib\n",
        "num_steps = 1000  # @param {type:\"number\"}\n",
        "example_viz_lib.run(agent='rainbow', game='SpaceInvaders', num_steps=num_steps,\n",
        "                    root_dir='/tmp/agent_viz', restore_ckpt='/tmp/tf_ckpt-199',\n",
        "                    use_legacy_checkpoint=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "sIO7Y1V_xJ50"
      },
      "outputs": [],
      "source": [
        "# @title Display the video\n",
        "import base64\n",
        "from IPython.display import HTML\n",
        "base_dir = '/tmp/agent_viz/agent_viz/SpaceInvaders/rainbow'\n",
        "video = open('{}/images/video.mp4'.format(base_dir), 'rb').read()\n",
        "encoded = base64.b64encode(video)\n",
        "HTML(data='''\u003cvideo alt=\"test\" controls\u003e\n",
        "                \u003csource src=\"data:video/mp4;base64,{0}\" type=\"video/mp4\" /\u003e\n",
        "             \u003c/video\u003e'''.format(encoded.decode('ascii')))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "ha3tGluP_rAL"
      },
      "outputs": [],
      "source": [
        "# @title Inspect individual video frames\n",
        "from IPython.display import Image\n",
        "frame_number = 100  # @param {type:\"slider\", min:0, max:999, step:1}\n",
        "image_file = '/tmp/agent_viz/agent_viz/SpaceInvaders/rainbow/images/frame_{:06d}.png'.format(frame_number)\n",
        "Image(image_file)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "agent_visualizer.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
